#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库迁移脚本
用于添加新增字段到现有数据库
"""
import sqlite3
import os

def migrate_database():
    """执行数据库迁移"""
    db_path = os.path.join(os.getcwd(), 'instance', 'coop_manager.db')
    
    if not os.path.exists(db_path):
        print("数据库文件不存在，将在应用启动时创建新的数据库")
        return
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        print("开始数据库迁移...")
        
        # 检查并添加customer表的新字段
        print("检查customer表...")
        cursor.execute("PRAGMA table_info(customer)")
        columns = [column[1] for column in cursor.fetchall()]
        
        new_customer_fields = [
            ('industry', 'VARCHAR(100)'),
            ('contact_person', 'VARCHAR(100)'),
            ('contact_phone', 'VARCHAR(20)'),
            ('contact_email', 'VARCHAR(100)'),
            ('address', 'TEXT'),
            ('description', 'TEXT'),
            ('contact_info', 'VARCHAR(200)')
        ]
        
        for field_name, field_type in new_customer_fields:
            if field_name not in columns:
                print(f"添加字段: customer.{field_name}")
                cursor.execute(f"ALTER TABLE customer ADD COLUMN {field_name} {field_type}")
        
        # 检查并添加staff表的新字段
        print("检查staff表...")
        cursor.execute("PRAGMA table_info(staff)")
        columns = [column[1] for column in cursor.fetchall()]
        
        new_staff_fields = [
            ('email', 'VARCHAR(100)'),
            ('contact_info', 'VARCHAR(200)')
        ]
        
        for field_name, field_type in new_staff_fields:
            if field_name not in columns:
                print(f"添加字段: staff.{field_name}")
                cursor.execute(f"ALTER TABLE staff ADD COLUMN {field_name} {field_type}")
        
        # 检查并添加negotiation表的新字段
        print("检查negotiation表...")
        cursor.execute("PRAGMA table_info(negotiation)")
        columns = [column[1] for column in cursor.fetchall()]
        
        new_negotiation_fields = [
            ('customer_id', 'INTEGER'),
            ('staff_id', 'INTEGER'),
            ('negotiation_date', 'DATE'),
            ('duration', 'INTEGER'),
            ('content', 'TEXT'),
            ('result', 'TEXT')
        ]
        
        for field_name, field_type in new_negotiation_fields:
            if field_name not in columns:
                print(f"添加字段: negotiation.{field_name}")
                cursor.execute(f"ALTER TABLE negotiation ADD COLUMN {field_name} {field_type}")
        
        # 检查并添加upload_exception表的新字段
        print("检查upload_exception表...")
        try:
            cursor.execute("PRAGMA table_info(upload_exception)")
            columns = [column[1] for column in cursor.fetchall()]
            
            new_exception_fields = [
                ('row_number', 'INTEGER'),
                ('error_message', 'TEXT'),
                ('raw_data', 'TEXT'),
                ('is_processed', 'BOOLEAN DEFAULT 0'),
                ('processed_at', 'DATETIME')
            ]
            
            for field_name, field_type in new_exception_fields:
                if field_name not in columns:
                    print(f"添加字段: upload_exception.{field_name}")
                    cursor.execute(f"ALTER TABLE upload_exception ADD COLUMN {field_name} {field_type}")
        except sqlite3.OperationalError:
            print("upload_exception表不存在，将由应用自动创建")
        
        # 检查并添加attendance表的新字段
        print("检查attendance表...")
        try:
            cursor.execute("PRAGMA table_info(attendance)")
            columns = [column[1] for column in cursor.fetchall()]
            
            new_attendance_fields = [
                ('staff_id', 'INTEGER'),
                ('date', 'DATE'),
                ('check_in_time', 'TIME'),
                ('check_out_time', 'TIME'),
                ('work_hours', 'FLOAT'),
                ('is_late', 'BOOLEAN DEFAULT 0'),
                ('is_early_leave', 'BOOLEAN DEFAULT 0'),
                ('status', 'VARCHAR(20) DEFAULT "正常"'),
                ('notes', 'TEXT')
            ]
            
            for field_name, field_type in new_attendance_fields:
                if field_name not in columns:
                    print(f"添加字段: attendance.{field_name}")
                    cursor.execute(f"ALTER TABLE attendance ADD COLUMN {field_name} {field_type}")
        except sqlite3.OperationalError:
            print("attendance表不存在，将由应用自动创建")
        
        # 检查并添加staff_project表的新字段
        print("检查staff_project表...")
        try:
            cursor.execute("PRAGMA table_info(staff_project)")
            columns = [column[1] for column in cursor.fetchall()]
            
            if 'role' not in columns:
                print("添加字段: staff_project.role")
                cursor.execute("ALTER TABLE staff_project ADD COLUMN role VARCHAR(50)")
        except sqlite3.OperationalError:
            print("staff_project表不存在，将由应用自动创建")
        
        # 检查并添加customer_staff表的新字段
        print("检查customer_staff表...")
        try:
            cursor.execute("PRAGMA table_info(customer_staff)")
            columns = [column[1] for column in cursor.fetchall()]
            
            if 'relationship_type' not in columns:
                print("添加字段: customer_staff.relationship_type")
                cursor.execute("ALTER TABLE customer_staff ADD COLUMN relationship_type VARCHAR(50)")
        except sqlite3.OperationalError:
            print("customer_staff表不存在，将由应用自动创建")
        
        conn.commit()
        print("数据库迁移完成！")
        
    except Exception as e:
        conn.rollback()
        print(f"迁移失败: {str(e)}")
        return False
    finally:
        conn.close()
    
    return True

if __name__ == '__main__':
    migrate_database()
